from sqlalchemy import create_engine
import pymysql
import pandas as pd
from Static import Static
from sqlalchemy.orm import sessionmaker


class BaseDao(object):
    def __init__(self):
        self.engine = create_engine(Static.DB_connect_address)
        self.Session = sessionmaker(bind=self.get_mysql_engine())

    # 批量无返回值操作,增删改
    def executeSQL_Statement(self, sqls):
        db = pymysql.connect(
            host=Static.MYSQL_host,
            user=Static.MYSQL_user,
            passwd=Static.MYSQL_pwd,
            port=Static.MYSQL_PORT,
            database=Static.MYSQL_DB,
            charset="utf8mb4",
        )
        cursor = db.cursor()
        try:
            for sql in sqls:
                cursor.execute(sql)
            # 提交事务
            db.commit()
        except Exception as e:
            print(e)
            db.rollback()
        finally:
            # 关闭游标
            cursor.close()
            # 关闭连接
            db.close()

    # 无返回值操作
    def executeNoResult(self, sql):
        self.executeSQL_Statement(sqls=[sql])

    # 查询
    def executeQuery(self, sql, GetRowNum=0):
        db = pymysql.connect(
            host=Static.MYSQL_host,
            user=Static.MYSQL_user,
            passwd=Static.MYSQL_pwd,
            port=Static.MYSQL_PORT,
            database=Static.MYSQL_DB,
            charset="utf8mb4",
        )
        cursor = db.cursor()
        try:
            cursor.execute(sql)
            if GetRowNum == 0:
                response = cursor.fetchall()
            elif GetRowNum == 1:
                response = cursor.fetchone()
            else:
                response = cursor.fetchmany(GetRowNum)
            data = pd.DataFrame(list(response))
        except Exception as e:
            print(e)
            data = pd.DataFrame()
        finally:
            # 关闭游标
            cursor.close()
            # 关闭连接
            db.close()
        return data

    def write2sql(self, df, db_name, write_type="append"):
        # session 会在 with 块结束时自动关闭
        with self.Session() as session:
            chunk_size = 1000  # 每次处理1000行
            df.to_sql(
                name=db_name,
                con=session.bind,
                if_exists=write_type,
                index=False,
                chunksize=chunk_size,
            )

    def readSQL(self, sql_query):
        with self.get_mysql_engine().connect() as con:
            df = pd.read_sql(sql=sql_query, con=con).fillna(0)
        return df

    def get_mysql_engine(self):
        return self.engine
